import numpy as np
import math
import trimesh
import pyrender
import PIL
import imageio
from io import BytesIO
from scipy.spatial.transform import Rotation as R



def renderMesh(mesh, width = 600, height = 600):
    extents = mesh.extents
    centroid = mesh.centroid
    scene = pyrender.Scene(ambient_light = [0.4,0.4,0.4, 0.4])
    scene.add(mesh)
    p0 = np.array([centroid[0] + 1.4*extents[0], centroid[1] + 1.4*extents[1], centroid[2] + 1.4*extents[2]])
    m = makeTranslationMatrix(p0)
    makeMatrixLookAt(m, centroid, up = np.array([0.0, 1.0, 0.0]))
    camera = pyrender.PerspectiveCamera(yfov=np.pi / 4.0, znear = np.amin(extents)*0.001 , zfar = np.amax(extents)*10 ,aspectRatio=1.0)
    camera_pose = m
    scene.add(camera, pose = camera_pose)
    d = np.array([1.4*extents[0], 1.4*extents[1], 1.4*extents[2]])
    camera2center = np.sqrt(d.dot(d))
    intensity = 1.5*math.pow(camera2center, 2.0)
    #print('intensity:', np.amax(extents)*120, 'camera distance:',np.sqrt(p0.dot(p0)))
    light = pyrender.PointLight(color = np.array([1.0,1.0,1.0]), intensity=intensity)
    light_pose = m
    scene.add(light, pose = light_pose)
    r = pyrender.OffscreenRenderer(width * 2, height * 2)
    color, depth = r.render(scene)
    im = PIL.Image.fromarray(color)
    im = im.resize((width, height))
    blob = BytesIO()
    im.save(blob, 'jpeg')
    blob.seek(0)
    return blob

def renderMeshAnimation(mesh, width = 200, height = 200):
    extents = mesh.extents
    centroid = mesh.centroid
    radius = np.array([np.amax(extents)]*3) * 1.4
    radius = np.sqrt(radius.dot(radius))
    scene = pyrender.Scene(ambient_light = [0.4,0.4,0.4, 1.0])
    scene.add(mesh)
    #p0 = np.array([centroid[0] + radius, centroid[1] , centroid[2]])
    #p0 = np.array([centroid[0] + 1.4*extents[0], centroid[1] + 1.4*extents[1], centroid[2] + 1.4*extents[2]])
    #m = makeTranslationMatrix(p0)
    #makeMatrixLookAt(m, centroid, up = np.array([0.0, 1.0, 0.0]))
    camera = pyrender.PerspectiveCamera(yfov=np.pi / 4.0, znear = np.amin(extents)*0.001 , zfar = np.amax(extents)*10 ,aspectRatio=1.0)
    #camera_pose = m
    p0 = np.array([radius, 0, 0])
    camera_node = scene.add(camera)
    camera2center = np.sqrt(p0.dot(p0))
    intensity = 1.5*math.pow(camera2center, 2.0)
    light = pyrender.PointLight(color = np.array([1.0,1.0,1.0]), intensity=intensity)
    light2 = pyrender.PointLight(color = np.array([1.0, 1.0, 1.0]), intensity=0)
    #light_pose = m
    light_node = scene.add(light)
    light_node2 = scene.add(light2)
    r = pyrender.OffscreenRenderer(2*width, 2*height)
    alpha = 0
    theta = 0
    up = np.array([0.0, 1.0, 0.0])
    index = 0
    step = 0.04
    count = np.floor(2.0*np.pi/step)
    array = np.zeros([int(count) + 1, width, height, 3])
    for theta in np.arange(0.0, 2.0*np.pi, step):
        #theta = alpha
        alpha = np.pi/3*np.sin(theta)
        theta2 = theta + np.pi/4
        p0 = np.array([centroid[0] + radius * np.cos(alpha) * np.cos(theta),\
centroid[1] + radius * np.sin(alpha), centroid[2] + radius * np.cos(alpha) * np.sin(theta)])
        p1 = np.array([centroid[0] + radius * np.cos(alpha) * np.cos(theta2),\
centroid[1] + radius * np.sin(alpha), centroid[2] + radius * np.cos(alpha) * np.sin(theta2)])
        m = makeTranslationMatrix(p0)
        makeMatrixLookAt(m, centroid, up = up)
        m2 = makeTranslationMatrix(p1)
        makeMatrixLookAt(m2, centroid, up = up)
        light_node.matrix = m
        camera_node.matrix = m
        light_node2.matrix = m2
        color, depth = r.render(scene)
        im = PIL.Image.fromarray(color)
        im = im.resize((width, height))
        newColor = np.array(im) 
        array[index] = newColor
        #im = PIL.Image.fromarray(color)
        #im.save(name + str(index) + ".jpeg")
        index = index + 1
    blob = BytesIO()
    imageio.mimwrite(blob, array,format = "gif" , fps = 30)
    blob.seek(0)
    return blob



#theta is scala value
#axis is a list of 1x3, etc [1,0,0]
def makeRotationMatrix(theta, axis):
    m4 = np.identity(4)
    s = np.sin(theta/2)
    c = np.cos(theta/2)
    quat = np.zeros([4])
    quat[0:3] = axis * s
    quat[3] = c
    r = R.from_quat(quat)
    m3 = r.as_matrix()
    m4[0:3, 0:3] = m3
    return m4

def makeTranslationMatrix(vec):
    m4 = np.identity(4)
    m4[0:3,3] = vec
    return m4

#vec must be a np.array
def makeMatrixLookAt(mat, vec, up = np.array([0,1.0,0])):
    vm = mat[0:3, 3]
    eye = vm - vec
    d = np.sqrt(eye.dot(eye))
    x = np.zeros([3])
    y = np.zeros([3])
    z = np.zeros([3])
    if d < 1e-16:
        z[2] = 1.0
    else:
        z = eye/d
    x = np.cross(up, z)
    if x.dot(x) < 1e-16:
        if 1- z[2] < 0.00001:
            z[2] += 0.00001
        else:
            z[0] += 0.00001
        z = z / np.sqrt(z.dot(z))
        x = np.cross(up, z)
    x = x / np.sqrt(x.dot(x))
    y = np.cross(z, x)
    mat[0:3, 0] = x
    mat[0:3, 1] = y
    mat[0:3, 2] = z
